from ollama import Client
import sys

def chat_with_ollama():
    # 初始化 Ollama 客户端
    client = Client(host='http://127.0.0.1:11434/')
    
    print("欢迎使用 Ollama 聊天程序！输入 'quit' 退出。")
    
    # 存储对话历史
    conversation = []
    
    try:
        while True:
            # 获取用户输入
            user_input = input("\n您: ")
            
            # 检查是否退出
            if user_input.lower() == 'quit':
                print("感谢使用，再见！")
                break
            
            # 将用户输入添加到对话历史
            conversation.append({
                'role': 'user',
                'content': user_input
            })
            
            try:
                # 调用 Ollama API
                response = client.chat(
                    model='qwen2.5:7b',  # 使用通义千问模型
                    messages=[
                        {
                            'role': 'system',
                            'content': '你是一个有帮助的AI助手，请用简洁明了的方式回答问题。'
                        },
                        *conversation  # 展开对话历史
                    ]
                )
                
                # 输出 AI 的回答
                ai_response = response['message']['content']
                print("\nAI:", ai_response)
                
                # 将 AI 的回答添加到对话历史
                conversation.append({
                    'role': 'assistant',
                    'content': ai_response
                })
                
            except Exception as e:
                print(f"\n错误: {str(e)}")
                
    except KeyboardInterrupt:
        print("\n程序已终止")
    
    except Exception as e:
        print(f"\n发生错误: {str(e)}")

if __name__ == "__main__":
    chat_with_ollama()
